package com.zdw.demo.controller;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.map.MapUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.coze.openapi.client.chat.model.ChatEvent;
import com.coze.openapi.client.chat.model.ChatEventType;
import com.coze.openapi.client.connversations.message.model.Message;
import com.coze.openapi.client.workflows.chat.WorkflowChatReq;
import com.coze.openapi.client.workflows.run.*;
import com.coze.openapi.client.workflows.run.model.WorkflowEvent;
import com.coze.openapi.client.workflows.run.model.WorkflowEventType;
import com.coze.openapi.client.workflows.run.model.WorkflowExecuteStatus;
import com.coze.openapi.client.workflows.run.model.WorkflowRunHistory;
import com.coze.openapi.service.auth.JWTOAuth;
import com.coze.openapi.service.auth.JWTOAuthClient;
import com.coze.openapi.service.config.Consts;
import com.coze.openapi.service.service.CozeAPI;
import com.zdw.common.constant.common.CommonConstants;
import com.zdw.common.exception.GlobalException;
import com.zdw.common.parent.JSONResult;
import com.zdw.demo.config.coze.CozeConfig;
import com.zdw.demo.dto.CozeResumeDto;
import com.zdw.demo.dto.CozeWorkflowDto;
import com.zdw.demo.vo.CozeWorkflowData;
import com.zdw.demo.vo.WorkflowRunHistoryVo;
import io.reactivex.Flowable;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;

import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Collections;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * Controller : 控制器类
 * 用于对请求的内容、响应的内容进行数据格式转换。
 */
@RestController                     //等于@ResponseBody(返回JSON格式的数据) + @Controller(定义为Controller接口类)
@RequestMapping(value = "coze")     //路径映射
@Slf4j                              //日志
public class CozeController {

    @Autowired
    private CozeConfig cozeConfig;

    private JWTOAuthClient oauth = null;

    private CozeAPI coze;

    /**
     * 初始化扣子oauth授权
     *
     * @throws Exception
     */
    private void oauth() throws Exception {
        String jwtOauthClientID           = cozeConfig.getCOZE_JWT_OAUTH_CLIENT_ID();
        String jwtOauthPrivateKeyFilePath = cozeConfig.getCOZE_JWT_OAUTH_PRIVATE_KEY_FILE_PATH();
        String jwtOauthPublicKeyID        = cozeConfig.getCOZE_JWT_OAUTH_PUBLIC_KEY_ID();

        /*
        The jwt oauth type requires using private to be able to issue a jwt token, and through
        the jwt token, apply for an access_token from the coze service. The sdk encapsulates
        this procedure, and only needs to use get_access_token to obtain the access_token under
        the jwt oauth process.
        Generate the authorization token
        The default ttl is 900s, and developers can customize the expiration time, which can be
        set up to 24 hours at most.
        * */
        String jwtOauthPrivateKey = new String(Files.readAllBytes(Paths.get(jwtOauthPrivateKeyFilePath)), StandardCharsets.UTF_8);
        if (ObjectUtil.isEmpty(oauth) || ObjectUtil.isEmpty(coze)) {
            oauth = new JWTOAuthClient.JWTOAuthBuilder()
                    .clientID(jwtOauthClientID)
                    .privateKey(jwtOauthPrivateKey)
                    .publicKey(jwtOauthPublicKeyID)
                    .baseURL(Consts.COZE_CN_BASE_URL)
                    .build();
            log.info("oauth初始化成功");

            coze = new CozeAPI.Builder()
                    .auth(JWTOAuth.builder().jwtClient(oauth).build())
                    .baseURL(Consts.COZE_CN_BASE_URL)
                    .build();
            log.info("coze初始化成功");
        }
    }

    /**
     * 启动扣子工作流
     * <p>
     * 限制说明:
     * 1)不支持通过此 API 执行包括以下节点的工作流。消息节点、开启了流式输出的结束节点、问答节点
     * 2)通过 API 方式执行工作流前，应确认此工作流已发布，执行从未发布过的工作流时会返回错误码 4200。创建并发布工作流的操作可参考使用工作流。
     * 3)异步运行的参数 is_async 仅限扣子专业版用户使用，否则调用此接口会报错 6003 Workflow execution with is_async=true is a premium feature available only to Coze Professional users.
     * 4)调用此 API 之前，应先在扣子平台中试运行此工作流，如果试运行时需要关联智能体，则调用此 API 执行工作流时，也需要指定智能体ID。通常情况下，执行存在数据库节点、变量节点等节点的工作流需要关联智能体
     *
     * @param cozeWorkflowDto 工作流入参
     */
    @PostMapping(value = "runCozeWorkflow")
    public JSONResult runCozeWorkflow(@RequestBody CozeWorkflowDto cozeWorkflowDto) {
        if (StrUtil.isEmpty(cozeWorkflowDto.getWorkflow_id())) {
            throw new GlobalException("工作流标识workflow_id不能为空");
        }
        try {
            oauth();
            // 创建工作流请求(同步/异步)
            RunWorkflowReq req = RunWorkflowReq.builder()
                    .workflowID(cozeWorkflowDto.getWorkflow_id())
                    .parameters(cozeWorkflowDto.getInput_map())
                    .isAsync(cozeWorkflowDto.is_async())
                    .build();
            // 调用应用工作流,需要增加appId
            if (StrUtil.isNotEmpty(cozeWorkflowDto.getApp_id())) {
                req = RunWorkflowReq.builder()
                        .workflowID(cozeWorkflowDto.getWorkflow_id())
                        .parameters(cozeWorkflowDto.getInput_map())
                        .isAsync(cozeWorkflowDto.is_async())
                        .appID(cozeWorkflowDto.getApp_id())
                        .build();
            }
            RunWorkflowResp runWorkflowResp = coze.workflows().runs().create(req);
            log.info("runWorkflowResp:{}", JSONUtil.toJsonStr(runWorkflowResp));
            boolean async = cozeWorkflowDto.is_async();
            if (runWorkflowResp.getCode() != 0) {
                return new JSONResult().markFail(CommonConstants.RESULT_FAIL_MSG, runWorkflowResp.getMsg());
            } else if (runWorkflowResp.getCode() == 0 && async) {
                return new JSONResult().markSuccess(CommonConstants.RESULT_SUCCESS_MSG, new CozeWorkflowData(runWorkflowResp.getExecuteID()));
            } else {
                CozeWorkflowData cozeWorkflowData = JSONUtil.toBean(runWorkflowResp.getData(), CozeWorkflowData.class);
                return new JSONResult().markSuccess(CommonConstants.RESULT_SUCCESS_MSG, cozeWorkflowData);
            }
        } catch (Exception e) {
            log.error("扣子工作流调用失败:{}", e.getMessage());
            throw new GlobalException("扣子工作流调用失败" + e.getMessage());
        }
    }

    /**
     * 查询扣子工作流异步执行结果
     * <p>
     * 限制说明:
     * 1)此接口对于每个工作流的 QPS 限频为 2，即针对某个工作流，每秒钟最多可查看 2 次执行结果。
     * 2)输出节点的输出数据最多保存 24 小时，结束节点为 7 天。
     *
     * @param workflow_id 工作流ID
     * @param execute_id  执行ID
     * @return
     */
    @GetMapping(value = "selectCozeAsyncWorkflow")
    public JSONResult selectCozeAsyncWorkflow(String workflow_id, String execute_id) {
        if (StrUtil.isEmpty(workflow_id)) {
            throw new GlobalException("工作流标识workflow_id不能为空");
        }
        if (StrUtil.isEmpty(execute_id)) {
            throw new GlobalException("执行标识execute_id不能为空");
        }
        boolean              isFinished           = false;
        int                  times                = 0;
        WorkflowRunHistoryVo workflowRunHistoryVo = null;
        WorkflowRunHistory   finalHistory         = null;
        try {
            oauth();
            while (!isFinished && times++ < 10) {
                RetrieveRunHistoryResp historyResp =
                        coze.workflows()
                                .runs()
                                .histories()
                                .retrieve(RetrieveRunHistoryReq.of(workflow_id, execute_id));
                log.info("historyResp:{}", JSONUtil.toJsonStr(historyResp));
                WorkflowRunHistory history = historyResp.getHistories().get(0);
                if (history.getExecuteStatus().equals(WorkflowExecuteStatus.FAIL)) {
                    log.info("Workflow run failed, reason:{},times:{}", history.getErrorMessage(), times);
                    isFinished = true;
                } else if (history.getExecuteStatus().equals(WorkflowExecuteStatus.RUNNING)) {
                    log.info("Workflow run is running,times:{}", times);
                    try {
                        TimeUnit.MILLISECONDS.sleep(500);
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                } else {
                    log.info("Workflow run success:{},times:{}", history.getOutput(), times);
                    isFinished = true;
                }
                if (isFinished || times == 10) {
                    finalHistory = history;
                    break;
                }
            }
            log.info("finalHistory:{}", JSONUtil.toJsonStr(finalHistory));
            JSONObject       outerJson        = JSONUtil.parseObj(finalHistory.getOutput());
            CozeWorkflowData cozeWorkflowData = JSONUtil.toBean(outerJson.getStr("Output"), CozeWorkflowData.class);
            workflowRunHistoryVo = new WorkflowRunHistoryVo(
                    finalHistory.getExecuteID(), finalHistory.getExecuteStatus().toString(), finalHistory.getCreateTime(),
                    finalHistory.getUpdateTime(), cozeWorkflowData, finalHistory.getErrorCode(), finalHistory.getErrorMessage()
            );
        } catch (Exception e) {
            log.error("查询扣子工作流运行状态失败:{}", e.getMessage());
            throw new GlobalException("查询扣子工作流运行状态失败" + e.getMessage());
        }
        return new JSONResult().markSuccess(CommonConstants.RESULT_SUCCESS_MSG, workflowRunHistoryVo);
    }

    /**
     * 启动扣子工作流(流式)
     * <p>
     * 限制说明:
     * 1)通过 API 方式执行工作流前，应确认此工作流已发布，执行从未发布过的工作流时会返回错误码 4200。创建并发布工作流的操作可参考使用工作流。
     * 2)调用此 API 之前，应先在扣子平台中试运行此工作流。
     * 如果试运行时需要关联智能体，则调用此 API 执行工作流时，也需要指定 bot_id。通常情况下，执行存在数据库节点、变量节点等节点的工作流需要关联智能体。
     * 执行应用中的工作流时，需要指定 app_id。
     * 请勿同时指定 bot_id 和 app_id，否则 API 会报错 4000，表示请求参数错误。
     * 3)此接口为同步接口，如果工作流整体或某些节点运行超时，智能体可能无法提供符合预期的回复。同步执行时，工作流整体超时时间限制可参考工作流使用限制。
     *
     * @param cozeWorkflowDto 工作流入参
     * 注意：前端需要使用支持流式响应的HTTP客户端，并逐块读取数据。
     */
    @PostMapping(value = "streamRunCozeWorkflow", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public StreamingResponseBody streamRunCozeWorkflow(@RequestBody CozeWorkflowDto cozeWorkflowDto) {
        if (StrUtil.isEmpty(cozeWorkflowDto.getWorkflow_id())) {
            throw new GlobalException("工作流标识workflow_id不能为空");
        }
        if (StrUtil.isNotEmpty(cozeWorkflowDto.getBot_id()) && StrUtil.isNotEmpty(cozeWorkflowDto.getApp_id())) {
            throw new GlobalException("请勿同时指定bot_id(智能体)和app_id(应用)");
        }
        Map<String, Object> inputMap = cozeWorkflowDto.getInput_map();
        if (CollectionUtil.isEmpty(inputMap) || StrUtil.isEmpty(MapUtil.getStr(inputMap, "input"))) {
            throw new GlobalException("参数input不能为空");
        }
        try {
            oauth();
            RunWorkflowReq req = RunWorkflowReq.builder()
                    .workflowID(cozeWorkflowDto.getWorkflow_id())
                    .readTimeout(300000)    //流式响应设置5分钟超时
                    .parameters(inputMap)
                    .build();
            Flowable<WorkflowEvent> flowable = coze.workflows().runs().stream(req);
            return out -> handleEvent(flowable, coze, out);
        } catch (Exception e) {
            log.error("扣子工作流调用失败:{}", e.getMessage());
            throw new GlobalException("扣子工作流(流式)调用失败" + e.getMessage());
        }
    }

    /**
     * 恢复运行工作流(流式响应)
     *
     * @param cozeResumeDto
     * @return
     */
    @PostMapping(value = "resumeCozeWorkflow", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public StreamingResponseBody resumeCozeWorkflow(@RequestBody CozeResumeDto cozeResumeDto) {
        if (StrUtil.isEmpty(cozeResumeDto.getWorkflow_id())) {
            throw new GlobalException("工作流标识workflow_id不能为空");
        }
        if (StrUtil.isEmpty(cozeResumeDto.getEvent_id())) {
            throw new GlobalException("中断事件ID event_id不能为空");
        }
        if (ObjectUtil.isEmpty(cozeResumeDto.getInterrupt_type())) {
            throw new GlobalException("中断类型interrupt_type不能为空");
        }
        if (StrUtil.isEmpty(cozeResumeDto.getResume_data())) {
            throw new GlobalException("问题回复resume_data不能为空");
        }
        try {
            oauth();
            ResumeRunReq req = ResumeRunReq.builder()
                    .workflowID(cozeResumeDto.getWorkflow_id())         //工作流ID
                    .eventID(cozeResumeDto.getEvent_id())               //中断事件ID
                    .interruptType(cozeResumeDto.getInterrupt_type())   //中断类型
                    .resumeData(cozeResumeDto.getResume_data())         //问题回复
                    .readTimeout(300000)                                //流式响应设置5分钟超时
                    .build();
            Flowable<WorkflowEvent> resume = coze.workflows().runs().resume(req);
            return out -> handleEvent(resume, coze, out);
        } catch (Exception e) {
            log.error("扣子工作流调用失败:{}", e.getMessage());
            throw new GlobalException("扣子工作流(流式)调用失败" + e.getMessage());
        }
    }

    /**
     * 处理工作流(流式)事件
     *
     * @param events
     * @param coze
     * @param out
     */
    private void handleEvent(Flowable<WorkflowEvent> events, CozeAPI coze, OutputStream out) {
        events.blockingForEach(
                event -> {
                    try {
                        if (event.getEvent().equals(WorkflowEventType.MESSAGE)) {
                            out.write(event.getMessage().getContent().getBytes());
                            log.info("工作流(流式响应)MESSAGE事件:{}", event.getMessage().getContent());
                        } else if (event.getEvent().equals(WorkflowEventType.ERROR)) {
                            out.write((event.getError().getErrorMessage()).getBytes());
                            log.info("工作流(流式响应)ERROR事件:{}", event.getError().getErrorMessage());
                        } else if (event.getEvent().equals(WorkflowEventType.DONE)) {
                            out.write((ObjectUtil.isEmpty(event.getMessage()) ? StrUtil.EMPTY : event.getMessage().getContent()).getBytes());
                            log.info("工作流(流式响应)MESSAGE事件:{}", ObjectUtil.isEmpty(event.getMessage()) ? StrUtil.EMPTY : event.getMessage().getContent());
                        } else if (event.getEvent().equals(WorkflowEventType.INTERRUPT)) {
                            out.write((ObjectUtil.isEmpty(event.getMessage()) ? StrUtil.EMPTY : event.getMessage().getContent()).getBytes());
                            log.info("工作流(流式响应)MESSAGE事件:{}", ObjectUtil.isEmpty(event.getMessage()) ? StrUtil.EMPTY : event.getMessage().getContent());
                        }
                        out.flush();
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                }
        );
        coze.shutdownExecutor();
    }

    /**
     * 执行对话流
     * <p>
     * 限制说明:
     * 1)通过 API 方式执行对话流前，应确认此对话流已发布，执行从未发布过的对话流时会返回错误码 4200。如果是扣子应用中的对话流，应先发布扣子应用为 API 服务；如果是空间资源库中的对话流，应先在资源库中发布对话流。
     * 2)此接口暂不支持异步运行。
     *
     * @param cozeWorkflowDto
     */
    @PostMapping(value = "streamRunCozeChatflow", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public StreamingResponseBody streamRunCozeChatflow(@RequestBody CozeWorkflowDto cozeWorkflowDto) {
        if (StrUtil.isEmpty(cozeWorkflowDto.getWorkflow_id())) {
            throw new GlobalException("工作流标识workflow_id不能为空");
        }
        if (StrUtil.isNotEmpty(cozeWorkflowDto.getBot_id()) && StrUtil.isNotEmpty(cozeWorkflowDto.getApp_id())) {
            throw new GlobalException("请勿同时指定bot_id(智能体)和app_id(应用)");
        }
        Map<String, Object> inputMap = cozeWorkflowDto.getInput_map();
        if (CollectionUtil.isEmpty(inputMap) || StrUtil.isEmpty(MapUtil.getStr(inputMap, "input"))) {
            throw new GlobalException("参数input不能为空");
        }
        try {
            oauth();

            WorkflowChatReq req = WorkflowChatReq.builder()
                    .workflowID(cozeWorkflowDto.getWorkflow_id())
                    .appID(cozeWorkflowDto.getApp_id())
                    .botID(cozeWorkflowDto.getBot_id())
                    .parameters(inputMap)
                    .additionalMessages(Collections.singletonList(Message.buildUserQuestionText(MapUtil.getStr(inputMap, "input"))))
                    .readTimeout(300000)    //流式响应设置5分钟超时
                    .build();
            Flowable<ChatEvent> events = coze.workflows().chat().stream(req);
            return out -> handleChatEvent(events, out);
        } catch (Exception e) {
            throw new GlobalException("扣子对话流调用失败" + e.getMessage());
        }
    }

    /**
     * 处理对话流(流式)事件
     *
     * @param events
     * @param out
     */
    private void handleChatEvent(Flowable<ChatEvent> events, OutputStream out) {
        events.blockingForEach(
                event -> {
                    if (ChatEventType.CONVERSATION_MESSAGE_DELTA.equals(event.getEvent())) {
                        log.info("event message: {}", event.getMessage().getContent());
                        out.write(event.getMessage().getContent().getBytes());
                        out.flush();
                    }
                });
    }

}
